use base64::Engine;
use itrustee_verifier::evidence::ITrusteeEvidence;
use itrustee_verifier::verifier::ItrusteePlugin;
use plugin_manager::{PluginError, ServiceHostFunctions, ServicePlugin};
use serde_json::Value;
use std::fs;
use std::path::PathBuf;

fn read_evidence_json() -> Value {
    let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    path.push("tests/data/report.json");
    let content = fs::read_to_string(path).expect("Failed to read input file");
    serde_json::from_str(&content).expect("Failed to parse input JSON")
}

fn create_default_host_functions() -> ServiceHostFunctions {
    ServiceHostFunctions {
        validate_cert_chain: Box::new(|_, _, _| Box::pin(async { true })),
        get_unmatched_measurements: Box::new(|_measured_values, _attester_type, _user_id| {
            Box::pin(async { Ok(Vec::new()) })
        }),
        query_configuration: |_key| None,
    }
}

/// Helper function to get the inner JSON from log_data and modify it
fn modify_inner_json<F>(right_report: &mut Value, modifier: F)
where
    F: FnOnce(&mut Value),
{
    if let Some(logs_array) = right_report.get_mut("logs").and_then(|l| l.as_array_mut()) {
        // Get the first log entry from the array
        if let Some(first_log) = logs_array.get_mut(0) {
            if let Some(log_data_b64) = first_log.get("log_data").and_then(|ld| ld.as_str()) {
                // Decode the base64 to get the inner JSON
                let raw = base64::engine::general_purpose::STANDARD.decode(log_data_b64).unwrap();
                let mut inner_json: Value = serde_json::from_slice(&raw).unwrap();

                // Apply the modifier function
                modifier(&mut inner_json);

                // Re-encode and set back
                let new_inner_str = serde_json::to_string(&inner_json).unwrap();
                let new_log_data = base64::engine::general_purpose::STANDARD.encode(new_inner_str);

                if let Some(log_obj) = first_log.as_object_mut() {
                    log_obj.insert("log_data".to_string(), serde_json::json!(new_log_data));
                }
            }
        }
    }
}

/// Helper function to modify akcert.sce_no_as fields
fn modify_sce_no_as<F>(right_report: &mut Value, modifier: F)
where
    F: FnOnce(&mut Value),
{
    modify_inner_json(right_report, |inner_json| {
        if let Some(akcert_obj) = inner_json.get_mut("akcert") {
            if let Some(sce_no_as) = akcert_obj.get_mut("sce_no_as") {
                modifier(sce_no_as);
            }
        }
    });
}

/// Helper function to modify akcert.sce_no_as.signature fields
fn modify_sce_no_as_signature<F>(right_report: &mut Value, modifier: F)
where
    F: FnOnce(&mut Value),
{
    modify_sce_no_as(right_report, |sce_no_as| {
        if let Some(signature) = sce_no_as.get_mut("signature") {
            modifier(signature);
        }
    });
}

/// Helper function to modify akcert.sce_no_as.payload fields
fn modify_sce_no_as_payload<F>(right_report: &mut Value, modifier: F)
where
    F: FnOnce(&mut Value),
{
    modify_sce_no_as(right_report, |sce_no_as| {
        if let Some(payload) = sce_no_as.get_mut("payload") {
            modifier(payload);
        }
    });
}

/// Helper function to modify akcert.sce_no_as.payload.ak_pub fields
fn modify_sce_no_as_ak_pub<F>(right_report: &mut Value, modifier: F)
where
    F: FnOnce(&mut Value),
{
    modify_sce_no_as_payload(right_report, |payload| {
        if let Some(ak_pub) = payload.get_mut("ak_pub") {
            modifier(ak_pub);
        }
    });
}

/// Correct NO AS scenario verification
#[tokio::test]
async fn test_itrustee_no_as_verify_success() {
    let report_json: Value = read_evidence_json();
    let mut evidence = ITrusteeEvidence::from_json_value(&report_json).expect("Failed to parse ITrusteeEvidence");

    let host_functions = create_default_host_functions();
    let plugin = ItrusteePlugin::new("itrustee".to_string(), host_functions);

    let ret = evidence.verify("", None, None, &plugin).await;
    assert!(ret.is_ok(), "NO AS verification should succeed");
}

/// Test invalid DRK certificate (base64 decode failure)
#[tokio::test]
async fn test_itrustee_no_as_invalid_drk_cert() {
    let mut report_json = read_evidence_json();

    // Modify drk_cert to be invalid base64
    modify_sce_no_as_signature(&mut report_json, |signature| {
        if let Some(drk_cert) = signature.get_mut("drk_cert") {
            *drk_cert = serde_json::json!("invalid_base64!@#$%");
        }
    });

    let mut evidence = ITrusteeEvidence::from_json_value(&report_json).unwrap();
    let host_functions = create_default_host_functions();
    let plugin = ItrusteePlugin::new("itrustee".to_string(), host_functions);

    let ret = evidence.verify("", None, None, &plugin).await;
    assert!(ret.is_err());

    match ret.err().unwrap() {
        PluginError::InternalError(msg) => {
            assert!(msg.contains("decode pem failed"));
        },
        _ => panic!("Expected InternalError for invalid DRK cert base64"),
    }
}

/// Test invalid DRK signature
#[tokio::test]
async fn test_itrustee_no_as_invalid_drk_sign() {
    let mut report_json = read_evidence_json();

    // Modify drk_sign to be invalid
    modify_sce_no_as_signature(&mut report_json, |signature| {
        if let Some(drk_sign) = signature.get_mut("drk_sign") {
            *drk_sign = serde_json::json!("invalid_signature");
        }
    });

    let mut evidence = ITrusteeEvidence::from_json_value(&report_json).unwrap();
    let host_functions = create_default_host_functions();
    let plugin = ItrusteePlugin::new("itrustee".to_string(), host_functions);

    let ret = evidence.verify("", None, None, &plugin).await;
    assert!(ret.is_err());

    // Should fail at signature verification step
    assert!(ret.is_err());
}


/// Test modified ak payload version causing signature verification failure
#[tokio::test]
async fn test_itrustee_no_as_modified_version_signature_fail() {
    let mut report_json = read_evidence_json();

    // Modify akcert.sce_no_as.payload.version
    modify_sce_no_as_payload(&mut report_json, |payload| {
        if let Some(version) = payload.get_mut("version") {
            *version = serde_json::json!("modified_version_2.0");
        }
    });

    let mut evidence = ITrusteeEvidence::from_json_value(&report_json).unwrap();
    let host_functions = create_default_host_functions();
    let plugin = ItrusteePlugin::new("itrustee".to_string(), host_functions);

    let ret = evidence.verify("", None, None, &plugin).await;
    assert!(ret.is_err());

    // Should fail at DRK signature verification step
    assert!(ret.is_err());
}

/// Test modified ak public key n field causing overall verification failure
#[tokio::test]
async fn test_itrustee_no_as_modified_ak_pub_n_verification_fail() {
    let mut report_json = read_evidence_json();

    // Modify akcert.sce_no_as.payload.ak_pub.n
    modify_sce_no_as_ak_pub(&mut report_json, |ak_pub| {
        if let Some(n) = ak_pub.get_mut("n") {
            *n = serde_json::json!("AQIDBAUGBwgJCgsMDQ4PEA==");
        }
    });

    let mut evidence = ITrusteeEvidence::from_json_value(&report_json).unwrap();
    let host_functions = create_default_host_functions();
    let plugin = ItrusteePlugin::new("itrustee".to_string(), host_functions);

    let ret = evidence.verify("", None, None, &plugin).await;
    assert!(ret.is_err());

    // Should fail at final AK signature verification step
    assert!(ret.is_err());
}

/// Test modified any payload field causing overall verification failure
#[tokio::test]
async fn test_itrustee_no_as_modified_payload_field_verification_fail() {
    let mut report_json = read_evidence_json();

    // Modify a field in the main payload
    modify_inner_json(&mut report_json, |inner_json| {
        if let Some(payload) = inner_json.get_mut("payload") {
            if let Some(timestamp) = payload.get_mut("timestamp") {
                *timestamp = serde_json::json!("modified_timestamp_1234567890");
            }
        }
    });

    let mut evidence = ITrusteeEvidence::from_json_value(&report_json).unwrap();
    let host_functions = create_default_host_functions();
    let plugin = ItrusteePlugin::new("itrustee".to_string(), host_functions);

    let ret = evidence.verify("", None, None, &plugin).await;
    assert!(ret.is_err());

    // Should fail at final AK signature verification step
    assert!(ret.is_err());
}

/// Test invalid final report signature
#[tokio::test]
async fn test_itrustee_no_as_invalid_final_signature() {
    let mut report_json = read_evidence_json();

    // Modify the final report signature to be invalid
    modify_inner_json(&mut report_json, |inner_json| {
        if let Some(report_sign_obj) = inner_json.get_mut("report_sign") {
            if let Some(sce_no_as) = report_sign_obj.get_mut("sce_no_as") {
                *sce_no_as = serde_json::json!("invalid_final_signature_base64");
            }
        }
    });

    let mut evidence = ITrusteeEvidence::from_json_value(&report_json).unwrap();
    let host_functions = create_default_host_functions();
    let plugin = ItrusteePlugin::new("itrustee".to_string(), host_functions);

    let ret = evidence.verify("", None, None, &plugin).await;
    assert!(ret.is_err());

    // Should fail at final signature verification step
    assert!(ret.is_err());
}

#[test]
fn test_get_sample_out() {
    let host_functions = create_default_host_functions();
    let plugin = ItrusteePlugin::new("itrustee".to_string(), host_functions);
    let sample = plugin.get_sample_output();

    let report_sign = sample.get("report_sign");
    assert!(report_sign.is_some(), "should have report sign");

    let akcert = sample.get("akcert");
    assert!(akcert.is_some(), "should have akcert");

    let payload = sample.get("payload");
    assert!(payload.is_some(), "should have a payload");
}